// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <algorithm>
#include <pollux/functions/lib/aggregates/compare.h>
#include <pollux/functions/lib/aggregates/min_max_by_aggregates_base.h>
#include <pollux/functions/lib/aggregates/value_set.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>
#include <pollux/type/floating_point_util.h>

namespace kumo::pollux::aggregate::prestosql {
    /// Returns true if the value in 'index' row of 'newComparisons' is strictly
    /// greater than or less than the value in the 'accumulator'.
    struct Comparator {
        template<
            bool greaterThan,
            typename T,
            typename TAccumulator,
            typename std::enable_if_t<std::is_same_v<T, TAccumulator>, int32_t>  = 0>
        static bool compare(
            TAccumulator *accumulator,
            const DecodedVector &newComparisons,
            vector_size_t index,
            bool isFirstValue) {
            if constexpr (functions::aggregate::isNumeric<T>()) {
                if (isFirstValue) {
                    return true;
                }
                if constexpr (greaterThan) {
                    if constexpr (std::is_floating_point_v<T>) {
                        return util::floating_point::NaNAwareGreaterThan<T>{}(
                            newComparisons.value_at<T>(index), *accumulator);
                    }
                    return newComparisons.value_at<T>(index) > *accumulator;
                } else {
                    if constexpr (std::is_floating_point_v<T>) {
                        return util::floating_point::NaNAwareLessThan<T>{}(
                            newComparisons.value_at<T>(index), *accumulator);
                    }
                    return newComparisons.value_at<T>(index) < *accumulator;
                }
            }
        }

        template<
            bool greaterThan,
            typename T,
            typename TAccumulator,
            typename std::enable_if_t<!std::is_same_v<T, TAccumulator>, int32_t>  = 0>
        static bool compare(
            TAccumulator *accumulator,
            const DecodedVector &newComparisons,
            vector_size_t index,
            bool /* isFirstValue */) {
            // SingleValueAccumulator::compare has the semantics of accumulator value
            // is less than vector value.
            static const CompareFlags kCompareFlags{
                true, // nullsFirst
                true, // ascending
                false, // equalsOnly
                CompareFlags::NullHandlingMode::kNullAsIndeterminate
            };
            if constexpr (greaterThan) {
                return !accumulator->hasValue() ||
                       functions::aggregate::compare(
                           accumulator, newComparisons, index, kCompareFlags) < 0;
            } else {
                return !accumulator->hasValue() ||
                       functions::aggregate::compare(
                           accumulator, newComparisons, index, kCompareFlags) > 0;
            }
        }
    };

    template<typename T>
    struct RawValueExtractor {
        using TRawValue = std::conditional_t<std::is_same_v<T, bool>, uint64_t, T>;

        static TRawValue *mutableRawValues(VectorPtr &values) {
            if constexpr (std::is_same_v<T, bool>) {
                return values->as<FlatVector<T> >()->template mutableRawValues<uint64_t>();
            } else {
                return values->as<FlatVector<T> >()->mutableRawValues();
            }
        }

        static void extract(TRawValue *rawValues, vector_size_t offset, T value) {
            if constexpr (std::is_same_v<T, bool>) {
                bits::setBit(rawValues, offset, value);
            } else {
                rawValues[offset] = value;
            }
        }
    };

    /// @tparam V Type of value.
    /// @tparam C Type of compare.
    /// @tparam Compare Type of comparator of std::pair<C, std::optional<V>>.
    template<typename V, typename C, typename Compare>
    struct MinMaxByNAccumulator {
        using TRawValue = typename RawValueExtractor<V>::TRawValue;
        using TRawComparison = typename RawValueExtractor<C>::TRawValue;

        int64_t n{0};

        using Pair = std::pair<C, std::optional<V> >;
        using Allocator = std::conditional_t<
            std::is_same_v<int128_t, V> || std::is_same_v<int128_t, C>,
            AlignedStlAllocator<Pair, sizeof(int128_t)>,
            StlAllocator<Pair> >;
        using Heap = std::vector<Pair, Allocator>;
        Heap heapValues;

        explicit MinMaxByNAccumulator(HashStringAllocator *allocator)
            : heapValues{Allocator(allocator)} {
        }

        int64_t getN() const {
            return n;
        }

        size_t size() const {
            return heapValues.size();
        }

        void checkAndSetN(DecodedVector &decodedN, vector_size_t row) {
            // Skip null N.
            if (decodedN.is_null_at(row)) {
                return;
            }

            const auto newN = decodedN.value_at<int64_t>(row);
            POLLUX_USER_CHECK_GT(
                newN, 0, "third argument of max_by/min_by must be a positive integer");

            if (n) {
                POLLUX_USER_CHECK_EQ(
                    newN,
                    n,
                    "third argument of max_by/min_by must be a constant for all rows in a group");
            } else {
                n = newN;
            }
        }

        void
        compareAndAdd(C comparison, std::optional<V> value, Compare &comparator) {
            if (heapValues.size() < n) {
                heapValues.push_back({comparison, value});
                std::push_heap(heapValues.begin(), heapValues.end(), comparator);
            } else {
                const auto &topPair = heapValues.front();
                if (comparator.compare(comparison, topPair)) {
                    std::pop_heap(heapValues.begin(), heapValues.end(), comparator);
                    heapValues.back() = std::make_pair(comparison, value);
                    std::push_heap(heapValues.begin(), heapValues.end(), comparator);
                }
            }
        }

        /// Extract all values from 'heapValues' into 'rawValues' and 'rawValueNulls'
        /// buffers. The heap remains unchanged after the call.
        void extractValues(
            TRawValue *rawValues,
            uint64_t *rawValueNulls,
            vector_size_t offset,
            Compare &comparator) {
            std::sort_heap(heapValues.begin(), heapValues.end(), comparator);
            // Add heap elements to rawValues in ascending order.
            for (int64_t i = 0; i < heapValues.size(); ++i) {
                const auto &pair = heapValues[i];
                const auto index = offset + i;
                const bool valueIsNull = !pair.second.has_value();
                bits::set_null(rawValueNulls, index, valueIsNull);
                if (!valueIsNull) {
                    RawValueExtractor<V>::extract(rawValues, index, pair.second.value());
                }
            }
            std::make_heap(heapValues.begin(), heapValues.end(), comparator);
        }

        /// Moves all pairs of (comparison, value) from 'heapValues' into
        /// 'rawComparisons', 'rawValues' and 'rawValueNulls' buffers. The heap
        /// remains unchanged after the call.
        void extractPairs(
            TRawComparison *rawComparisons,
            TRawValue *rawValues,
            uint64_t *rawValueNulls,
            vector_size_t offset,
            Compare &comparator) {
            std::sort_heap(heapValues.begin(), heapValues.end(), comparator);
            // Add heap elements to rawComparisons and rawValues in ascending order.
            for (int64_t i = 0; i < heapValues.size(); ++i) {
                const auto &pair = heapValues[i];
                const auto index = offset + i;

                RawValueExtractor<C>::extract(rawComparisons, index, pair.first);

                const bool valueIsNull = !pair.second.has_value();
                bits::set_null(rawValueNulls, index, valueIsNull);
                if (!valueIsNull) {
                    RawValueExtractor<V>::extract(rawValues, index, pair.second.value());
                }
            }
            std::make_heap(heapValues.begin(), heapValues.end(), comparator);
        }
    };

    template<typename V, typename C, typename Compare>
    struct Extractor {
        using TRawValue = typename RawValueExtractor<V>::TRawValue;
        using TRawComparison = typename RawValueExtractor<C>::TRawValue;

        TRawValue *rawValues;
        uint64_t *rawValueNulls;

        explicit Extractor(VectorPtr &values) {
            rawValues = RawValueExtractor<V>::mutableRawValues(values);
            rawValueNulls = values->mutable_raw_nulls();
        }

        void extractValues(
            MinMaxByNAccumulator<V, C, Compare> *accumulator,
            vector_size_t offset,
            Compare &comparator) {
            accumulator->extractValues(rawValues, rawValueNulls, offset, comparator);
        }

        void extractPairs(
            MinMaxByNAccumulator<V, C, Compare> *accumulator,
            TRawComparison *rawComparisons,
            vector_size_t offset,
            Compare &comparator) {
            accumulator->extractPairs(
                rawComparisons, rawValues, rawValueNulls, offset, comparator);
        }
    };

    template<typename V, typename C, typename Compare>
    struct MinMaxByNStringViewAccumulator {
        using BaseType = MinMaxByNAccumulator<V, C, Compare>;
        BaseType base;
        ValueSet valueSet;

        explicit MinMaxByNStringViewAccumulator(HashStringAllocator *allocator)
            : base{allocator}, valueSet{allocator} {
        }

        ~MinMaxByNStringViewAccumulator() {
            for (auto i = 0; i < base.heapValues.size(); ++i) {
                freePair(base.heapValues[i]);
            }
        }

        int64_t getN() const {
            return base.n;
        }

        size_t size() const {
            return base.size();
        }

        void checkAndSetN(DecodedVector &decodedN, vector_size_t row) {
            return base.checkAndSetN(decodedN, row);
        }

        void
        compareAndAdd(C comparison, std::optional<V> value, Compare &comparator) {
            if (base.heapValues.size() < base.n) {
                addToAccumulator(comparison, value, comparator);
            } else {
                const auto &topPair = base.heapValues.front();
                if (comparator.compare(comparison, topPair)) {
                    std::pop_heap(
                        base.heapValues.begin(), base.heapValues.end(), comparator);
                    freePair(base.heapValues.back());
                    base.heapValues.pop_back();
                    addToAccumulator(comparison, value, comparator);
                }
            }
        }

        /// Extract all values from 'heapValues' into 'rawValues' and 'rawValueNulls'
        /// buffers. The heap remains unchanged after the call.
        void extractValues(
            FlatVector<V> &values,
            vector_size_t offset,
            Compare &comparator) {
            std::sort_heap(base.heapValues.begin(), base.heapValues.end(), comparator);
            // Add heap elements to values in ascending order.
            for (int64_t i = 0; i < base.heapValues.size(); ++i) {
                const auto &pair = base.heapValues[i];
                extractValue(pair, values, offset + i);
            }
            std::make_heap(base.heapValues.begin(), base.heapValues.end(), comparator);
        }

        /// Moves all pairs of (comparison, value) from 'heapValues' into
        /// 'rawComparisons', 'rawValues' and 'rawValueNulls' buffers. The heap
        /// remains unchanged after the call.
        void extractPairs(
            FlatVector<C> &compares,
            FlatVector<V> &values,
            vector_size_t offset,
            Compare &comparator) {
            std::sort_heap(base.heapValues.begin(), base.heapValues.end(), comparator);
            // Add heap elements to compares and values in ascending order.
            for (int64_t i = 0; i < base.heapValues.size(); ++i) {
                const auto &pair = base.heapValues[i];
                const auto index = offset + i;

                extractCompare(pair, compares, index);
                extractValue(pair, values, index);
            }
            std::make_heap(base.heapValues.begin(), base.heapValues.end(), comparator);
        }

    private:
        using Pair = typename MinMaxByNAccumulator<V, C, Compare>::Pair;

        std::optional<StringView> writeString(
            const std::optional<StringView> &value) {
            if (!value.has_value()) {
                return std::nullopt;
            }
            return valueSet.write(*value);
        }

        void
        addToAccumulator(C comparison, std::optional<V> value, Compare &comparator) {
            if constexpr (
                std::is_same_v<V, StringView> && std::is_same_v<C, StringView>) {
                base.heapValues.push_back(
                    std::make_pair(valueSet.write(comparison), writeString(value)));
            } else if constexpr (std::is_same_v<V, StringView>) {
                base.heapValues.push_back(std::make_pair(comparison, writeString(value)));
            } else {
                static_assert(
                    std::is_same_v<C, StringView>,
                    "At least one of V and C must be StringView.");
                base.heapValues.push_back(
                    std::make_pair(valueSet.write(comparison), value));
            }
            std::push_heap(base.heapValues.begin(), base.heapValues.end(), comparator);
        }

        void freePair(typename BaseType::Heap::const_reference pair) {
            if constexpr (std::is_same_v<C, StringView>) {
                valueSet.free(pair.first);
            }
            if constexpr (std::is_same_v<V, StringView>) {
                if (pair.second.has_value()) {
                    valueSet.free(*pair.second);
                }
            }
        }

        void
        extractValue(const Pair &pair, FlatVector<V> &values, vector_size_t index) {
            const bool valueIsNull = !pair.second.has_value();
            values.set_null(index, valueIsNull);
            if (!valueIsNull) {
                values.set(index, pair.second.value());
            }
        }

        void extractCompare(
            const Pair &pair,
            FlatVector<C> &compares,
            vector_size_t index) {
            compares.set_null(index, false);
            compares.set(index, pair.first);
        }
    };

    template<typename V, typename C, typename Compare>
    struct StringViewExtractor {
        using TRawComparison = typename RawValueExtractor<C>::TRawValue;
        FlatVector<V> &values;
        FlatVector<C> *compares;

        StringViewExtractor(VectorPtr &_values, BaseVector *_compares)
            : values{*_values->as_flat_vector<V>()},
              compares{_compares ? _compares->as_flat_vector<C>() : nullptr} {
        }

        void extractValues(
            MinMaxByNStringViewAccumulator<V, C, Compare> *accumulator,
            vector_size_t offset,
            Compare &comparator) {
            accumulator->extractValues(values, offset, comparator);
        }

        void extractPairs(
            MinMaxByNStringViewAccumulator<V, C, Compare> *accumulator,
            vector_size_t offset,
            Compare &comparator) {
            POLLUX_DCHECK_NOT_NULL(compares);
            accumulator->extractPairs(*compares, values, offset, comparator);
        }
    };

    /// @tparam C Type of compare.
    /// @tparam Compare Type of comparator of
    /// std::pair<C, std::optional<HashStringAllocator::Position>>.
    template<typename C, typename Compare>
    struct MinMaxByNComplexTypeAccumulator {
        using BaseType =
        MinMaxByNAccumulator<HashStringAllocator::Position, C, Compare>;
        BaseType base;
        ValueSet valueSet;

        explicit MinMaxByNComplexTypeAccumulator(HashStringAllocator *allocator)
            : base{allocator}, valueSet{allocator} {
        }

        ~MinMaxByNComplexTypeAccumulator() {
            for (auto i = 0; i < base.heapValues.size(); ++i) {
                freePair(base.heapValues[i]);
            }
        }

        int64_t getN() const {
            return base.n;
        }

        size_t size() const {
            return base.size();
        }

        void checkAndSetN(DecodedVector &decodedN, vector_size_t row) {
            return base.checkAndSetN(decodedN, row);
        }

        void compareAndAdd(
            C comparison,
            DecodedVector &decoded,
            vector_size_t index,
            Compare &comparator) {
            if (base.heapValues.size() < base.n) {
                auto position = writeComplex(decoded, index);
                addToAccumulator(comparison, position, comparator);
            } else {
                const auto &topPair = base.heapValues.front();
                if (comparator.compare(comparison, topPair)) {
                    std::pop_heap(
                        base.heapValues.begin(), base.heapValues.end(), comparator);
                    auto position = writeComplex(decoded, index);
                    freePair(base.heapValues.back());
                    base.heapValues.pop_back();
                    addToAccumulator(comparison, position, comparator);
                }
            }
        }

        /// Extract all values from 'heapValues' into 'rawValues' and 'rawValueNulls'
        /// buffers. The heap remains unchanged after the call.
        void
        extractValues(BaseVector &values, vector_size_t offset, Compare &comparator) {
            std::sort_heap(base.heapValues.begin(), base.heapValues.end(), comparator);
            // Add heap elements to values in ascending order.
            for (int64_t i = 0; i < base.heapValues.size(); ++i) {
                const auto &pair = base.heapValues[i];
                extractValue(pair, values, offset + i);
            }
            std::make_heap(base.heapValues.begin(), base.heapValues.end(), comparator);
        }

        /// Moves all pairs of (comparison, value) from 'heapValues' into
        /// 'rawComparisons', 'rawValues' and 'rawValueNulls' buffers. The heap
        /// remains unchanged after the call.
        void extractPairs(
            FlatVector<C> &compares,
            BaseVector &values,
            vector_size_t offset,
            Compare &comparator) {
            std::sort_heap(base.heapValues.begin(), base.heapValues.end(), comparator);
            // Add heap elements to compares and values in ascending order.
            for (int64_t i = 0; i < base.heapValues.size(); ++i) {
                const auto &pair = base.heapValues[i];
                const auto index = offset + i;

                extractCompare(pair, compares, index);
                extractValue(pair, values, index);
            }
            std::make_heap(base.heapValues.begin(), base.heapValues.end(), comparator);
        }

    private:
        using V = HashStringAllocator::Position;
        using Pair = typename MinMaxByNAccumulator<V, C, Compare>::Pair;

        void addToAccumulator(
            C comparison,
            const std::optional<V> &position,
            Compare &comparator) {
            if constexpr (std::is_same_v<C, StringView>) {
                base.heapValues.push_back(
                    std::make_pair(valueSet.write(comparison), position));
            } else {
                base.heapValues.push_back(std::make_pair(comparison, position));
            }
            std::push_heap(base.heapValues.begin(), base.heapValues.end(), comparator);
        }

        std::optional<HashStringAllocator::Position> writeComplex(
            DecodedVector &decoded,
            vector_size_t index) {
            if (decoded.is_null_at(index)) {
                return std::nullopt;
            }
            HashStringAllocator::Position position;
            valueSet.write(*decoded.base(), decoded.index(index), position);
            return position;
        }

        void freePair(typename BaseType::Heap::const_reference pair) {
            if constexpr (std::is_same_v<C, StringView>) {
                valueSet.free(pair.first);
            }
            if (pair.second.has_value()) {
                valueSet.free(pair.second->header);
            }
        }

        void extractValue(const Pair &pair, BaseVector &values, vector_size_t index) {
            const bool valueIsNull = !pair.second.has_value();
            values.set_null(index, valueIsNull);
            if (!valueIsNull) {
                auto position = pair.second.value();
                valueSet.read(&values, index, position.header);
            }
        }

        void extractCompare(
            const Pair &pair,
            FlatVector<C> &compares,
            vector_size_t index) {
            compares.set_null(index, false);
            compares.set(index, pair.first);
        }
    };

    template<typename C, typename Compare>
    struct ComplexTypeExtractor {
        BaseVector &values;
        FlatVector<C> *compares;

        explicit ComplexTypeExtractor(VectorPtr &_values, BaseVector *_compares)
            : values{*_values},
              compares{_compares ? _compares->as_flat_vector<C>() : nullptr} {
        }

        void extractValues(
            MinMaxByNComplexTypeAccumulator<C, Compare> *accumulator,
            vector_size_t offset,
            Compare &comparator) {
            accumulator->extractValues(values, offset, comparator);
        }

        void extractPairs(
            MinMaxByNComplexTypeAccumulator<C, Compare> *accumulator,
            vector_size_t offset,
            Compare &comparator) {
            POLLUX_DCHECK_NOT_NULL(compares);
            accumulator->extractPairs(*compares, values, offset, comparator);
        }
    };

    template<typename V, typename C>
    struct Less {
        using Pair = std::pair<C, std::optional<V> >;

        bool operator()(const Pair &lhs, const Pair &rhs) {
            if constexpr (std::is_floating_point_v<C>) {
                return util::floating_point::NaNAwareLessThan<C>{}(lhs.first, rhs.first);
            }
            return lhs.first < rhs.first;
        }

        bool compare(C lhs, const Pair &rhs) {
            if constexpr (std::is_floating_point_v<C>) {
                return util::floating_point::NaNAwareLessThan<C>{}(lhs, rhs.first);
            }
            return lhs < rhs.first;
        }
    };

    template<typename V, typename C>
    struct Greater {
        using Pair = std::pair<C, std::optional<V> >;

        bool operator()(const Pair &lhs, const Pair &rhs) {
            if constexpr (std::is_floating_point_v<C>) {
                return util::floating_point::NaNAwareGreaterThan<C>{}(
                    lhs.first, rhs.first);
            }
            return lhs.first > rhs.first;
        }

        bool compare(C lhs, const Pair &rhs) {
            if constexpr (std::is_floating_point_v<C>) {
                return util::floating_point::NaNAwareGreaterThan<C>{}(lhs, rhs.first);
            }
            return lhs > rhs.first;
        }
    };

    template<typename V, typename C, typename Compare>
    struct MinMaxByNAccumulatorTypeTraits {
        using AccumulatorType = MinMaxByNAccumulator<V, C, Compare>;
        using ExtractorType = Extractor<V, C, Compare>;
    };

    // Use MinMaxByNStringViewAccumulator and StringViewExtractor if at least one of
    // V and C is StringView. (Unless V is ComplexType.)
    template<typename C, typename Compare>
    struct MinMaxByNAccumulatorTypeTraits<StringView, C, Compare> {
        using AccumulatorType =
        MinMaxByNStringViewAccumulator<StringView, C, Compare>;
        using ExtractorType = StringViewExtractor<StringView, C, Compare>;
    };

    template<typename Compare>
    struct MinMaxByNAccumulatorTypeTraits<StringView, StringView, Compare> {
        using AccumulatorType =
        MinMaxByNStringViewAccumulator<StringView, StringView, Compare>;
        using ExtractorType = StringViewExtractor<StringView, StringView, Compare>;
    };

    template<typename V, typename Compare>
    struct MinMaxByNAccumulatorTypeTraits<V, StringView, Compare> {
        using AccumulatorType =
        MinMaxByNStringViewAccumulator<V, StringView, Compare>;
        using ExtractorType = StringViewExtractor<V, StringView, Compare>;
    };

    // Use MinMaxByNComplexTypeAccumulator and ComplexTypeExtractor if V is
    // ComplexType.
    template<typename C, typename Compare>
    struct MinMaxByNAccumulatorTypeTraits<ComplexType, C, Compare> {
        using AccumulatorType = MinMaxByNComplexTypeAccumulator<C, Compare>;
        using ExtractorType = ComplexTypeExtractor<C, Compare>;
    };

    template<typename Compare>
    struct MinMaxByNAccumulatorTypeTraits<ComplexType, StringView, Compare> {
        using AccumulatorType = MinMaxByNComplexTypeAccumulator<StringView, Compare>;
        using ExtractorType = ComplexTypeExtractor<StringView, Compare>;
    };

    namespace {
        std::pair<vector_size_t *, vector_size_t *> rawOffsetAndSizes(
            ArrayVector &array_vector) {
            return {
                array_vector.offsets()->asMutable<vector_size_t>(),
                array_vector.sizes()->asMutable<vector_size_t>()
            };
        }
    } // namespace

    template<typename V, typename C, typename Compare>
    class MinMaxByNAggregate : public exec::Aggregate {
    public:
        explicit MinMaxByNAggregate(TypePtr resultType)
            : exec::Aggregate(resultType) {
        }

        using AccumulatorType =
        typename MinMaxByNAccumulatorTypeTraits<V, C, Compare>::AccumulatorType;
        using ExtractorType =
        typename MinMaxByNAccumulatorTypeTraits<V, C, Compare>::ExtractorType;
        using TRawComparison = typename RawValueExtractor<C>::TRawValue;

        int32_t accumulatorFixedWidthSize() const override {
            return sizeof(AccumulatorType);
        }

        void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
        override {
            auto valuesArray = (*result)->as<ArrayVector>();
            valuesArray->resize(numGroups);

            const auto numValues =
                    countValuesAndSetResultNulls(groups, numGroups, *result);

            auto values = valuesArray->elements();
            values->resize(numValues);

            std::optional<ExtractorType> extractor;
            if constexpr (
                std::is_same_v<V, StringView> || std::is_same_v<C, StringView> ||
                std::is_same_v<V, ComplexType>) {
                extractor.emplace(values, nullptr);
            } else {
                extractor.emplace(values);
            }

            auto [rawOffsets, rawSizes] = rawOffsetAndSizes(*valuesArray);

            vector_size_t offset = 0;
            for (auto i = 0; i < numGroups; ++i) {
                auto *group = groups[i];

                if (!isNull(group)) {
                    auto *accumulator = value(group);
                    const vector_size_t size = accumulator->size();

                    rawOffsets[i] = offset;
                    rawSizes[i] = size;

                    extractor->extractValues(accumulator, offset, comparator_);

                    offset += size;
                }
            }
        }

        void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
        override {
            auto row_vector = (*result)->as<RowVector>();
            row_vector->resize(numGroups);

            auto nVector = row_vector->childAt(0);
            auto comparisonArray = row_vector->childAt(1)->as<ArrayVector>();
            auto valueArray = row_vector->childAt(2)->as<ArrayVector>();

            auto *rawNs = nVector->as<FlatVector<int64_t> >()->mutableRawValues();

            const auto numValues =
                    countValuesAndSetResultNulls(groups, numGroups, *result);

            auto values = valueArray->elements();
            auto comparisons = comparisonArray->elements();

            values->resize(numValues);
            comparisons->resize(numValues);

            std::optional<ExtractorType> extractor;
            if constexpr (
                std::is_same_v<V, StringView> || std::is_same_v<C, StringView> ||
                std::is_same_v<V, ComplexType>) {
                extractor.emplace(values, comparisons.get());
            } else {
                extractor.emplace(values);
            }

            TRawComparison *rawComparisons =
                    RawValueExtractor<C>::mutableRawValues(comparisons);

            auto [rawValueOffsets, rawValueSizes] = rawOffsetAndSizes(*valueArray);
            auto [rawComparisonOffsets, rawComparisonSizes] =
                    rawOffsetAndSizes(*comparisonArray);

            vector_size_t offset = 0;
            for (auto i = 0; i < numGroups; ++i) {
                auto *group = groups[i];

                if (!isNull(group)) {
                    auto *accumulator = value(group);
                    const auto size = accumulator->size();

                    rawNs[i] = accumulator->getN();

                    rawValueOffsets[i] = offset;
                    rawValueSizes[i] = size;

                    rawComparisonOffsets[i] = offset;
                    rawComparisonSizes[i] = size;

                    if constexpr (
                        std::is_same_v<V, StringView> || std::is_same_v<C, StringView> ||
                        std::is_same_v<V, ComplexType>) {
                        extractor->extractPairs(accumulator, offset, comparator_);
                    } else {
                        extractor->extractPairs(
                            accumulator, rawComparisons, offset, comparator_);
                    }

                    offset += size;
                }
            }
        }

        void addRawInput(
            char **groups,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool /*unused*/) override {
            decodedValue_.decode(*args[0], rows);
            decodedComparison_.decode(*args[1], rows);
            decodedN_.decode(*args[2], rows);

            rows.applyToSelected([&](vector_size_t i) {
                if (decodedComparison_.is_null_at(i) || decodedN_.is_null_at(i)) {
                    return;
                }

                auto *group = groups[i];

                auto *accumulator = value(group);
                accumulator->checkAndSetN(decodedN_, i);

                addRawInput(group, i);
            });
        }

        void addIntermediateResults(
            char **groups,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool /*mayPushdown*/) override {
            auto results = decodeIntermediateResults(args[0], rows);

            rows.applyToSelected([&](vector_size_t i) {
                if (!decodedIntermediates_.is_null_at(i)) {
                    addIntermediateResults(groups[i], i, results);
                }
            });
        }

        void addSingleGroupRawInput(
            char *group,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool /*unused*/) override {
            decodedValue_.decode(*args[0], rows);
            decodedComparison_.decode(*args[1], rows);

            auto *accumulator = value(group);
            validateN(args[2], rows, accumulator);

            rows.applyToSelected([&](vector_size_t i) {
                if (!decodedComparison_.is_null_at(i) && !decodedN_.is_null_at(i)) {
                    addRawInput(group, i);
                }
            });
        }

        void addSingleGroupIntermediateResults(
            char *group,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool /*mayPushdown*/) override {
            auto results = decodeIntermediateResults(args[0], rows);

            rows.applyToSelected([&](vector_size_t i) {
                if (!decodedIntermediates_.is_null_at(i)) {
                    addIntermediateResults(group, i, results);
                }
            });
        }

    protected:
        void initializeNewGroupsInternal(
            char **groups,
            melon::Range<const vector_size_t *> indices) override {
            exec::Aggregate::setAllNulls(groups, indices);
            for (const vector_size_t i: indices) {
                auto group = groups[i];
                new(group + offset_) AccumulatorType(allocator_);
            }
        }

        void destroyInternal(melon::Range<char **> groups) override {
            destroyAccumulators<AccumulatorType>(groups);
        }

    private:
        inline AccumulatorType *value(char *group) {
            return reinterpret_cast<AccumulatorType *>(group + Aggregate::offset_);
        }

        static std::optional<V> optionalValue(
            const DecodedVector &decoded,
            vector_size_t index) {
            std::optional<V> value;
            if (!decoded.is_null_at(index)) {
                value = decoded.value_at<V>(index);
            }

            return value;
        }

        static std::optional<V> optionalValue(
            const FlatVector<V> &vector,
            vector_size_t index) {
            std::optional<V> value;
            if (!vector.is_null_at(index)) {
                value = vector.value_at(index);
            }

            return value;
        }

        void addRawInput(char *group, vector_size_t index) {
            clearNull(group);

            auto *accumulator = value(group);

            const auto comparison = decodedComparison_.value_at<C>(index);
            if constexpr (std::is_same_v<V, ComplexType>) {
                accumulator->compareAndAdd(comparison, decodedValue_, index, comparator_);
            } else {
                const auto value = optionalValue(decodedValue_, index);
                accumulator->compareAndAdd(comparison, value, comparator_);
            }
        }

        struct IntermediateResult {
            const ArrayVector *valueArray;
            // Used for complex types.
            DecodedVector values;
            // Used for primitive types.
            const FlatVector<V> *flatValues;
            const ArrayVector *comparisonArray;
            const FlatVector<C> *comparisons;
        };

        void addIntermediateResults(
            char *group,
            vector_size_t index,
            IntermediateResult &result) {
            clearNull(group);

            auto *accumulator = value(group);

            const auto decodedIndex = decodedIntermediates_.index(index);

            accumulator->checkAndSetN(decodedN_, decodedIndex);

            const auto *valueArray = result.valueArray;
            const auto *values = result.flatValues;
            const auto *comparisonArray = result.comparisonArray;
            const auto *comparisons = result.comparisons;

            const auto numValues = valueArray->sizeAt(decodedIndex);
            const auto valueOffset = valueArray->offsetAt(decodedIndex);
            const auto comparisonOffset = comparisonArray->offsetAt(decodedIndex);
            for (auto i = 0; i < numValues; ++i) {
                const auto comparison = comparisons->value_at(comparisonOffset + i);
                if constexpr (std::is_same_v<V, ComplexType>) {
                    accumulator->compareAndAdd(
                        comparison, result.values, valueOffset + i, comparator_);
                } else {
                    const auto value = optionalValue(*values, valueOffset + i);
                    accumulator->compareAndAdd(comparison, value, comparator_);
                }
            }
        }

        IntermediateResult decodeIntermediateResults(
            const VectorPtr &arg,
            const SelectivityVector &rows) {
            decodedIntermediates_.decode(*arg, rows);

            auto baseRowVector =
                    dynamic_cast<const RowVector *>(decodedIntermediates_.base());

            decodedN_.decode(*baseRowVector->childAt(0), rows);
            decodedComparison_.decode(*baseRowVector->childAt(1), rows);
            decodedValue_.decode(*baseRowVector->childAt(2), rows);

            IntermediateResult result{};
            result.valueArray = decodedValue_.base()->template as<ArrayVector>();
            result.comparisonArray =
                    decodedComparison_.base()->template as<ArrayVector>();

            if constexpr (std::is_same_v<V, ComplexType>) {
                result.values.decode(*result.valueArray->elements());
            } else {
                result.flatValues =
                        result.valueArray->elements()->template as<FlatVector<V> >();
            }
            result.comparisons =
                    result.comparisonArray->elements()->template as<FlatVector<C> >();

            return result;
        }

        /// Return total number of values in all accumulators of specified 'groups'.
        /// Set null flags in 'result'.
        vector_size_t countValuesAndSetResultNulls(
            char **groups,
            int32_t numGroups,
            VectorPtr &result) {
            vector_size_t numValues = 0;

            uint64_t *rawNulls = getRawNulls(result.get());

            for (auto i = 0; i < numGroups; ++i) {
                auto *group = groups[i];
                auto *accumulator = value(group);

                if (isNull(group)) {
                    result->set_null(i, true);
                } else {
                    clearNull(rawNulls, i);
                    numValues += accumulator->size();
                }
            }

            return numValues;
        }

        void validateN(
            const VectorPtr &arg,
            const SelectivityVector &rows,
            AccumulatorType *accumulator) {
            decodedN_.decode(*arg, rows);
            if (decodedN_.isConstantMapping()) {
                accumulator->checkAndSetN(decodedN_, rows.begin());
            } else {
                rows.applyToSelected(
                    [&](auto row) { accumulator->checkAndSetN(decodedN_, row); });
            }
        }

        Compare comparator_;
        DecodedVector decodedValue_;
        DecodedVector decodedComparison_;
        DecodedVector decodedN_;
        DecodedVector decodedIntermediates_;
    };

    template<template <typename U, typename V> class NAggregate, typename W>
    std::unique_ptr<exec::Aggregate> createNArg(
        TypePtr resultType,
        TypePtr compareType,
        const std::string &errorMessage) {
        switch (compareType->kind()) {
            case TypeKind::BOOLEAN:
                return std::make_unique<NAggregate<W, bool> >(resultType);
            case TypeKind::TINYINT:
                return std::make_unique<NAggregate<W, int8_t> >(resultType);
            case TypeKind::SMALLINT:
                return std::make_unique<NAggregate<W, int16_t> >(resultType);
            case TypeKind::INTEGER:
                return std::make_unique<NAggregate<W, int32_t> >(resultType);
            case TypeKind::BIGINT:
                return std::make_unique<NAggregate<W, int64_t> >(resultType);
            case TypeKind::HUGEINT:
                return std::make_unique<NAggregate<W, int128_t> >(resultType);
            case TypeKind::REAL:
                return std::make_unique<NAggregate<W, float> >(resultType);
            case TypeKind::DOUBLE:
                return std::make_unique<NAggregate<W, double> >(resultType);
            case TypeKind::VARCHAR:
                return std::make_unique<NAggregate<W, StringView> >(resultType);
            case TypeKind::TIMESTAMP:
                return std::make_unique<NAggregate<W, Timestamp> >(resultType);
            default:
                POLLUX_FAIL("{}", errorMessage);
                return nullptr;
        }
    }

    template<template <typename U, typename V> class NAggregate>
    std::unique_ptr<exec::Aggregate> createNArg(
        TypePtr resultType,
        TypePtr valueType,
        TypePtr compareType,
        const std::string &errorMessage) {
        switch (valueType->kind()) {
            case TypeKind::BOOLEAN:
                return createNArg<NAggregate, bool>(
                    resultType, compareType, errorMessage);
            case TypeKind::TINYINT:
                return createNArg<NAggregate, int8_t>(
                    resultType, compareType, errorMessage);
            case TypeKind::SMALLINT:
                return createNArg<NAggregate, int16_t>(
                    resultType, compareType, errorMessage);
            case TypeKind::INTEGER:
                return createNArg<NAggregate, int32_t>(
                    resultType, compareType, errorMessage);
            case TypeKind::BIGINT:
                return createNArg<NAggregate, int64_t>(
                    resultType, compareType, errorMessage);
            case TypeKind::HUGEINT:
                return createNArg<NAggregate, int128_t>(
                    resultType, compareType, errorMessage);
            case TypeKind::REAL:
                return createNArg<NAggregate, float>(
                    resultType, compareType, errorMessage);
            case TypeKind::DOUBLE:
                return createNArg<NAggregate, double>(
                    resultType, compareType, errorMessage);
            case TypeKind::VARBINARY:
                [[fallthrough]];
            case TypeKind::VARCHAR:
                return createNArg<NAggregate, StringView>(
                    resultType, compareType, errorMessage);
            case TypeKind::TIMESTAMP:
                return createNArg<NAggregate, Timestamp>(
                    resultType, compareType, errorMessage);
            case TypeKind::ARRAY:
                [[fallthrough]];
            case TypeKind::MAP:
                [[fallthrough]];
            case TypeKind::ROW:
                return createNArg<NAggregate, ComplexType>(
                    resultType, compareType, errorMessage);
            default:
                POLLUX_FAIL(errorMessage);
        }
    }

    inline std::string toString(const std::vector<TypePtr> &types) {
        std::ostringstream out;
        for (auto i = 0; i < types.size(); ++i) {
            if (i > 0) {
                out << ", ";
            }
            out << types[i]->toString();
        }
        return out.str();
    }

    template<
        template <
            typename U,
            typename V,
            bool B1,
            class C,
            bool compareTypeUsesCustomComparison>
        class Aggregate,
        bool isMaxFunc,
        template <typename U, typename V>
        class NAggregate>
    exec::AggregateRegistrationResult registerMinMaxBy(
        const std::string &name,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures;
        // V, C -> row(V, C) -> V.
        signatures.push_back(exec::AggregateFunctionSignatureBuilder()
            .typeVariable("V")
            .orderableTypeVariable("C")
            .returnType("V")
            .intermediateType("row(V,C)")
            .argumentType("V")
            .argumentType("C")
            .build());
        // Add signatures for 3-arg version of min_by/max_by.
        const std::vector<std::string> supportedCompareTypes = {
            "boolean",
            "tinyint",
            "smallint",
            "integer",
            "bigint",
            "hugeint",
            "real",
            "double",
            "varchar",
            "date",
            "timestamp"
        };

        // Add support for all value types to 3-arg version of the aggregate.
        for (const auto &compareType: supportedCompareTypes) {
            // V, C, bigint -> row(bigint, array(C), array(V)) -> array(V).
            signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                .typeVariable("V")
                .returnType("array(V)")
                .intermediateType(fmt::format(
                    "row(bigint,array({}),array(V))", compareType))
                .argumentType("V")
                .argumentType(compareType)
                .argumentType("bigint")
                .build());
        }
        signatures.push_back(
            exec::AggregateFunctionSignatureBuilder()
            .integerVariable("a_precision")
            .integerVariable("a_scale")
            .typeVariable("V")
            .returnType("array(V)")
            .intermediateType(
                "row(bigint,array(DECIMAL(a_precision, a_scale)),array(V))")
            .argumentType("V")
            .argumentType("DECIMAL(a_precision, a_scale)")
            .argumentType("bigint")
            .build());
        return exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step step,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                const std::string errorMessage = fmt::format(
                    "Unknown input types for {} ({}) aggregation: {}",
                    name,
                    mapAggregationStepToName(step),
                    toString(argTypes));

                const bool nAgg = (argTypes.size() == 3) ||
                                  (argTypes.size() == 1 && argTypes[0]->size() == 3);

                if (nAgg) {
                    return createNArg<NAggregate>(
                        resultType, argTypes[0], argTypes[1], errorMessage);
                } else {
                    if (argTypes[1]->providesCustomComparison()) {
                        return functions::aggregate::
                                create<Aggregate, Comparator, isMaxFunc, true>(
                                    resultType, argTypes[0], argTypes[1], errorMessage, true);
                    } else {
                        return functions::aggregate::
                                create<Aggregate, Comparator, isMaxFunc, false>(
                                    resultType, argTypes[0], argTypes[1], errorMessage, true);
                    }
                }
            },
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
